import torch_geometric.utils
from torch_geometric.data import HeteroData
import torch
import matplotlib.pyplot as plt
import networkx as nx
import torch_geometric.transforms as T
from tokenizer_tools import get_sentence_vector
from memory import memory_usage


class MyData(object):
    def __init__(self, alarms_df,batch_size=400):
        num_features_alarm = 768
        num_features_host = 768
        num_bussiness_trees = 768
        data = HeteroData()
        # 获取行数
        idx_dict = self.get_item_indexes(alarms_df[0:100])
        # print(idx_dict.keys())
        # print(len(idx_dict["contents_keys"]))

        data["alarm"].node_id = torch.arange(0, len(idx_dict["contents_keys"]))
        data['alarm'].x = self.get_feature_vector(idx_dict["contents_keys"], num_features_alarm)

        data["host"].node_id = torch.arange(0, len(idx_dict["ips_keys"]))
        data['host'].x = self.get_feature_vector(idx_dict["ips_keys"], num_features_host)

        data["bussiness_tree"].node_id = torch.arange(0, len(idx_dict["product_labels_keys"]))
        data['bussiness_tree'].x = self.get_feature_vector(idx_dict["product_labels_keys"], num_bussiness_trees)

        data['alarm', 'on', 'host'].edge_index = idx_dict["content2ip_idx"].to(torch.int64)
        data['alarm', 'to', 'bussiness_tree'].edge_index = idx_dict["content2label_idx"].to(torch.int64)
        data['host', 'belongsto', 'bussiness_tree'].edge_index = idx_dict["ip2label_idx"].to(torch.int64)
        # data = data.pin_memory()
        # data = data.to('cuda:0', non_blocking=True)
        alarm_dict = {i: idx_dict["contents_keys"][i] for i in range(len(idx_dict["contents_keys"]))}
        host_dict = {i: idx_dict["ips_keys"][i] for i in range(len(idx_dict["ips_keys"]))}
        bussiness_tree_dict = {i: idx_dict["product_labels_keys"][i] for i in
                               range(len(idx_dict["product_labels_keys"]))}

        all_dict = {i: idx_dict["contents_keys"][i] for i in range(len(idx_dict["contents_keys"]))}
        for i in range(len(idx_dict["ips_keys"])):
            all_dict[i + len(idx_dict["contents_keys"])] = idx_dict["ips_keys"][i]
        for i in range(len(idx_dict["product_labels_keys"])):
            all_dict[i + len(idx_dict["contents_keys"]) + len(idx_dict["ips_keys"])] = idx_dict["product_labels_keys"][
                i]

        data["alarm_dict"] = alarm_dict
        data["host_dict"] = host_dict
        data["bussiness_tree_dict"] = bussiness_tree_dict
        data["all_dict"] = all_dict

        self._data = data

    @property
    def data(self):
        return self._data

    def get_item_indexes(self, df):
        df["combined"] = df.apply(
            lambda row: "|".join([row["ip"], row["content"], row["idc_name"], row["product_labels"]]), axis=1)

        ori_product_labels = df["product_labels"].tolist()
        product_labels = []
        for l in ori_product_labels:
            product_labels.extend(l.split(","))

        contents = df["combined"].tolist()
        product_labels = list(set(product_labels))
        ips = list(set(df["ip"].tolist()))

        idx_dict = {
            "contents_keys": contents,
            "product_labels_keys": product_labels,
            "ips_keys": ips,
        }
        # make feature
        # content_list = []
        # for i in contents:
        #     alarm_explain = get_sentence_vector(i)
        #     content_list.append(alarm_explain)
        # contents_feature= torch.stack(content_list, dim=0).view([-1, num_features_alarm])
        # idx_dict.update({"contents_feature":contents_feature})
        #
        # product_label_list = []
        # for i in product_labels:
        #     alarm_explain = get_sentence_vector(i)
        #     product_label_list.append(alarm_explain)
        # product_labels_feature= torch.stack(product_label_list, dim=0).view([-1, num_features_alarm])
        # idx_dict.update({"product_labels_feature":product_labels_feature})

        # ip_list = []
        # for i in ips:
        #     alarm_explain = get_sentence_vector(i)
        #     ip_list.append(alarm_explain)
        # ips_feature= torch.stack(ip_list, dim=0).view([-1, num_features_alarm])
        # idx_dict.update({"ips_feature":ips_feature})

        # make edge indexes 1
        alarm2ip_list = []
        for i in range(len(contents)):
            alarm2ip_list.append([i, ips.index(contents[i].split("|")[0])])
        # print(torch.IntTensor(alarm2ip_list).T)
        alarm2ip_idx = torch.IntTensor(alarm2ip_list).T
        idx_dict.update({"content2ip_idx": alarm2ip_idx})

        # print(product_labels)
        # print(contents)
        # make edge indexes 2
        alarm2product_label_list = []
        for i in range(len(contents)):
            labels = contents[i].split("|")[-1].split(",")
            for l in labels:
                alarm2product_label_list.append([i, product_labels.index(l)])
        # print(torch.IntTensor(alarm2ip_list).T)
        alarm2label_idx = torch.IntTensor(alarm2product_label_list).T
        idx_dict.update({"content2label_idx": alarm2label_idx})

        # make edge indexes 3
        ip2product_label_list = []
        for i in range(len(contents)):
            ip = contents[i].split("|")[0]
            labels = contents[i].split("|")[-1].split(",")
            for l in labels:
                if [ips.index(ip), product_labels.index(l)] not in ip2product_label_list:
                    ip2product_label_list.append([ips.index(ip), product_labels.index(l)])
        ip2label_idx = torch.IntTensor(ip2product_label_list).T
        idx_dict.update({"ip2label_idx": ip2label_idx})

        # print(idx_dict)
        return idx_dict

    def get_feature_vector(self, list_A, num_features_alarm):
        f_list = []

        for i in list_A:
            memory_usage()
            alarm_explain = get_sentence_vector(i)
            f_list.append(alarm_explain)
        return torch.stack(f_list, dim=0).view([-1, num_features_alarm])


# # 将 PyTorch Geometric 图转换为 NetworkX 图
# g = torch_geometric.utils.to_networkx(data, to_undirected=False)
# print(data)
# node_colors = ['red' for i in range(10)]+['green' for i in range(10)]+['blue' for i in range(10)]  # 对应每个节点的颜色
# # 绘制图形
# # nx.draw(g, pos=nx.spring_layout(g),with_labels=True,node_color=node_colors)
# # # 保存为 PNG 图像
# nx.draw_networkx_labels(g, nx.spring_layout(g),  font_color="red")
# nx.draw_networkx_edge_labels(g, nx.spring_layout(g),
#                              font_color="blue")  # no edge_labels parameter, default is showing all attributes of edges
# plt.show()
# plt.savefig("output.png")


# data['paper', 'cites', 'paper'].edge_attr = ...  # [num_edges_cites, num_features_cites]
# data['author', 'writes', 'paper'].edge_attr = ...  # [num_edges_writes, num_features_writes]
# data['author', 'affiliated_with', 'institution'].edge_attr = ...  # [num_edges_affiliated, num_features_affiliated]
# data['paper', 'has_topic', 'field_of_study'].edge_attr = ...  # [num_edges_topic, num_features_topic]

# print(data.x_dict.keys())
# print(data.edge_index_dict.keys())
# print(data.metadata())

